package org.example.process;

import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessAllWindowFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.table.planner.expressions.In;
import org.apache.flink.util.Collector;
import org.apache.kafka.common.protocol.types.Field;
import org.example.data.WaterSensor;
import org.example.function.WaterSensorMapFunction;

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 滑动时间窗口，实现TopN
 */
public class KeyProcessFuncTopNDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<String> source = env.socketTextStream("localhost", 9999);
        SingleOutputStreamOperator<WaterSensor> waterSensorStream = source.map(new WaterSensorMapFunction());

        WatermarkStrategy<WaterSensor> waterSensorWatermarkStrategy = WatermarkStrategy
                //指定乱序的Watermark生成，延时。
                .<WaterSensor>forBoundedOutOfOrderness(Duration.ofSeconds(3))
                .withTimestampAssigner((WaterSensor waterSensor, long l) -> {
                    //返回的时间戳 毫秒
                    return waterSensor.getTs() * 1000;
                });

        SingleOutputStreamOperator<WaterSensor> sensorDS = waterSensorStream.assignTimestampsAndWatermarks(waterSensorWatermarkStrategy);
        //思路二：使用KeyedProcessFunction实现
        //最近10s=窗口长度，每5s输出=滑动步长
        //1.按VC进行分组、开窗、聚合（增量累加，全量打标签）
        //因为开窗聚合后就是普通的DataStream，没有窗口信息，需要自己打上窗标志windowEnd
        KeyedStream<WaterSensor, Integer> sensorKeyDS = sensorDS.keyBy(WaterSensor::getVc);
        SingleOutputStreamOperator<Tuple3<Integer, Integer, Long>> windowAgg = sensorKeyDS.window(SlidingEventTimeWindows.of(Time.seconds(10), Time.seconds(5)))
                .aggregate(new VcCountAggFunc(), new WindowResult());
        //2.按照窗口结束时间进行KeyBy，保证同一个窗口时间范围的数据，放到一起去，之后再排序
        windowAgg.keyBy(tup -> tup.f2)
                .process(new TopN(2))
                .print();


        env.execute("PartitionDemo");
    }

    public static class VcCountAggFunc implements AggregateFunction<WaterSensor, Integer, Integer> {
        @Override
        public Integer createAccumulator() {
            return 0;
        }

        @Override
        public Integer add(WaterSensor value, Integer accumulator) {
            return accumulator + 1;
        }

        @Override
        public Integer getResult(Integer accumulator) {
            return accumulator;
        }

        @Override
        public Integer merge(Integer a, Integer b) {
            return null;
        }
    }

    /**
     * 泛型说明：
     * 第1个输入类型等于增量函数的输出，Count
     * 第2个输出的类型：Tuple3<vc, count, windowEnd>，带上窗口的结束时长
     * 第3个Key的类型
     * 第4个窗口的类型
     */
    public static class WindowResult extends ProcessWindowFunction<Integer, Tuple3<Integer, Integer, Long>, Integer, TimeWindow> {
        private static final long serialVersionUID = 1L;

        @Override
        public void process(Integer key, ProcessWindowFunction<Integer, Tuple3<Integer, Integer, Long>, Integer, TimeWindow>.Context context, Iterable<Integer> elements, Collector<Tuple3<Integer, Integer, Long>> out) throws Exception {
            //迭代器里面只有一条数据，next一次即可
            out.collect(Tuple3.of(key, elements.iterator().next(), context.window().getEnd()));
        }
    }

    /**
     * 处理Top N
     */
    public static class TopN extends KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String> {
        //存不同窗口的统计结果，Key为windowEnd，值为数据列表
        private Map<Long, List<Tuple3<Integer, Integer, Long>>> map = new HashMap<>();
        //取得Top几
        private Integer threshold;

        public TopN(Integer threshold) {
            this.threshold = threshold;
        }

        @Override
        public void processElement(Tuple3<Integer, Integer, Long> value, KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String>.Context ctx, Collector<String> out) throws Exception {
            //来一条处理一次,要排序需要等数据先到齐==> 暂存起来不同的窗口分开存
            //1。数据累加
            map.computeIfAbsent(ctx.getCurrentKey(), k -> new ArrayList<>()).add(value);
            //2.注册一个定时器,windowEnd+1ms即可
            //同一个窗口范围，应该同时输出，只不过一条一条处理processElement，只需要延时1ms
            ctx.timerService().registerEventTimeTimer(ctx.getCurrentKey() + 1);
        }


        @Override
        public void onTimer(long timestamp, KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String>.OnTimerContext ctx, Collector<String> out) throws Exception {
            //定时器触发，输出统计结果
            List<Tuple3<Integer, Integer, Long>> list = map.get(ctx.getCurrentKey());
            list.sort((o1, o2) -> o2.f1 - o1.f1);
            StringBuilder outStr = new StringBuilder();
            outStr.append("<===========================>\n");
            Integer i = 0;
            for (Tuple3<Integer, Integer, Long> o : list) {
                if (i++ >= threshold) break;
                outStr.append("Top: " + i + " 的vc: " + o.f0 + " 值为:" + o.f1 + "窗口结束时间：" + DateFormatUtils.format(ctx.getCurrentKey(), "yyyy-MM-dd HH:mm:ss") + "\n");
            }
            outStr.append("<===========================>\n");
            out.collect(outStr.toString());
            //清除数据
            map.remove(ctx.getCurrentKey());
        }
    }
}